{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "20d2d16a-7f75-42e9-9745-e08ee8ccb309",
   "metadata": {},
   "source": [
    "# 导入环境"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "12ddcb7d-b41c-4e68-bcfd-dc192e39d19e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset\n",
    "import pandas as pd\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq, TrainingArguments, Trainer, GenerationConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "53c06dd1-8a7f-4bb9-a0d6-a96a4ccd3898",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 读取json数据集文件\n",
    "df = pd.read_json('dataset/huanhuan.json')\n",
    "ds = Dataset.from_pandas(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9b146696-f7c1-45e9-9160-6609427cd182",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'instruction': ['小姐，别的秀女都在求中选，唯有咱们小姐想被撂牌子，菩萨一定记得真真儿的——',\n",
       "  '这个温太医啊，也是古怪，谁不知太医不得皇命不能为皇族以外的人请脉诊病，他倒好，十天半月便往咱们府里跑。',\n",
       "  '嬛妹妹，刚刚我去府上请脉，听甄伯母说你来这里进香了。'],\n",
       " 'input': ['', '', ''],\n",
       " 'output': ['嘘——都说许愿说破是不灵的。', '你们俩话太多了，我该和温太医要一剂药，好好治治你们。', '出来走走，也是散心。']}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds[:3] # 展示前三组数据"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac65caa6-2e8c-4794-aadf-2a59e6a0bde7",
   "metadata": {},
   "source": [
    "# 处理数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bbd74576-8329-4f8c-80b2-8959b7a25214",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = 'autodl-tmp/LLM-Research/gemma-3-4b-it'\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False, trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "db3cf8c9-8899-4da9-8f48-2dc5813420fe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<bos><start_of_turn>user\n",
      "You are a helpful assistant.\n",
      "\n",
      "你好呀<end_of_turn>\n",
      "<start_of_turn>model\n",
      "有什么可以帮你的？<end_of_turn>\n",
      "<start_of_turn>model\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# 使用tokenizer构建messages并打印， 查看chat_template的输出格式\n",
    "messages = [\n",
    "            {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
    "            {\"role\": \"user\", \"content\": '你好呀'},\n",
    "            {\"role\": \"assistant\", \"content\": '有什么可以帮你的？'}\n",
    "            ]\n",
    "# 使用chat_template将messages格式化并打印\n",
    "print(tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "325b08dc-7d5d-4c4d-9dfd-af097ee1647f",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = '现在你要扮演皇帝身边的女人--甄嬛'\n",
    "\n",
    "def process_func(example):\n",
    "    MAX_LENGTH = 384    # 分词器会将一个中文字切分为多个token，因此需要放开一些最大长度，保证数据的完整性\n",
    "    input_ids, attention_mask, labels = [], [], []\n",
    "    # 构建指令部分的输入, 可参考上面的输出格式进行调整和补充\n",
    "    instruction = tokenizer(\n",
    "        f\"<s><|im_start|>system\\n{system_prompt}<|im_end|>\\n\" \n",
    "        f\"<|im_start|>user\\n{example['instruction'] + example['input']}<|im_end|>\\n\"  \n",
    "        f\"<|im_start|>assistant\\n\",  \n",
    "        add_special_tokens=False   \n",
    "    )\n",
    "    # 构建模型回复部分的输入\n",
    "    response = tokenizer(\n",
    "        f\"{example['output']}\",\n",
    "        add_special_tokens=False \n",
    "    )\n",
    "    # 拼接指令和回复部分的 input_ids\n",
    "    input_ids = instruction[\"input_ids\"] + response[\"input_ids\"] + [tokenizer.pad_token_id]\n",
    "    # 拼接指令和回复部分的 attention_mask\n",
    "    attention_mask = instruction[\"attention_mask\"] + response[\"attention_mask\"] + [1]  # 因为 EOS token 也需要关注，所以补充为 1\n",
    "    # 构建标签\n",
    "    # 对于指令部分，使用 -100 忽略其损失计算；对于回复部分，保留其 input_ids 作为标签\n",
    "    labels = [-100] * len(instruction[\"input_ids\"]) + response[\"input_ids\"] + [tokenizer.pad_token_id]  \n",
    "    # 如果总长度超过最大长度，进行截断\n",
    "    if len(input_ids) > MAX_LENGTH: \n",
    "        input_ids = input_ids[:MAX_LENGTH]\n",
    "        attention_mask = attention_mask[:MAX_LENGTH]\n",
    "        labels = labels[:MAX_LENGTH]\n",
    "    return {\n",
    "        \"input_ids\": input_ids,\n",
    "        \"attention_mask\": attention_mask,\n",
    "        \"labels\": labels\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1429b7d8-a41a-4ea3-b26c-d45cdb2fdd2b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9d6209c945154d919cb78b38d50c99c5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/3729 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['input_ids', 'attention_mask', 'labels'],\n",
       "    num_rows: 3729\n",
       "})"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenized_id = ds.map(process_func, remove_columns=ds.column_names)\n",
    "tokenized_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5a1e11ce-e9ab-4d73-a6eb-6ab22969d69c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s><|im_start|>system\n",
      "现在你要扮演皇帝身边的女人--甄嬛<|im_end|>\n",
      "<|im_start|>user\n",
      "这个温太医啊，也是古怪，谁不知太医不得皇命不能为皇族以外的人请脉诊病，他倒好，十天半月便往咱们府里跑。<|im_end|>\n",
      "<|im_start|>assistant\n",
      "你们俩话太多了，我该和温太医要一剂药，好好治治你们。<pad>\n"
     ]
    }
   ],
   "source": [
    "# 解码输入\n",
    "print(tokenizer.decode(tokenized_id[1]['input_ids']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1bb9fc43-69f9-4ab4-86e1-51c95a4eac0e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "你们俩话太多了，我该和温太医要一剂药，好好治治你们。<pad>\n"
     ]
    }
   ],
   "source": [
    "# 解码标签, 过滤掉-100\n",
    "print(tokenizer.decode(list(filter(lambda x: x != -100, tokenized_id[1][\"labels\"]))))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37a0fc7f-7adc-4926-b55f-3cb7dd932be6",
   "metadata": {},
   "source": [
    "# 创建模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2e34e258-9f38-4b49-9716-2b1d358c3aa1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4a36c1b2c3b345cb9a09290a10564aec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Gemma3ForConditionalGeneration(\n",
       "  (vision_tower): SiglipVisionModel(\n",
       "    (vision_model): SiglipVisionTransformer(\n",
       "      (embeddings): SiglipVisionEmbeddings(\n",
       "        (patch_embedding): Conv2d(3, 1152, kernel_size=(14, 14), stride=(14, 14), padding=valid)\n",
       "        (position_embedding): Embedding(4096, 1152)\n",
       "      )\n",
       "      (encoder): SiglipEncoder(\n",
       "        (layers): ModuleList(\n",
       "          (0-26): 27 x SiglipEncoderLayer(\n",
       "            (self_attn): SiglipSdpaAttention(\n",
       "              (k_proj): Linear(in_features=1152, out_features=1152, bias=True)\n",
       "              (v_proj): Linear(in_features=1152, out_features=1152, bias=True)\n",
       "              (q_proj): Linear(in_features=1152, out_features=1152, bias=True)\n",
       "              (out_proj): Linear(in_features=1152, out_features=1152, bias=True)\n",
       "            )\n",
       "            (layer_norm1): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)\n",
       "            (mlp): SiglipMLP(\n",
       "              (activation_fn): PytorchGELUTanh()\n",
       "              (fc1): Linear(in_features=1152, out_features=4304, bias=True)\n",
       "              (fc2): Linear(in_features=4304, out_features=1152, bias=True)\n",
       "            )\n",
       "            (layer_norm2): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (post_layernorm): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)\n",
       "    )\n",
       "  )\n",
       "  (multi_modal_projector): Gemma3MultiModalProjector(\n",
       "    (mm_soft_emb_norm): Gemma3RMSNorm((1152,), eps=1e-06)\n",
       "    (avg_pool): AvgPool2d(kernel_size=4, stride=4, padding=0)\n",
       "  )\n",
       "  (language_model): Gemma3ForCausalLM(\n",
       "    (model): Gemma3TextModel(\n",
       "      (embed_tokens): Gemma3TextScaledWordEmbedding(262208, 2560, padding_idx=0)\n",
       "      (layers): ModuleList(\n",
       "        (0-33): 34 x Gemma3DecoderLayer(\n",
       "          (self_attn): Gemma3Attention(\n",
       "            (q_proj): Linear(in_features=2560, out_features=2048, bias=False)\n",
       "            (k_proj): Linear(in_features=2560, out_features=1024, bias=False)\n",
       "            (v_proj): Linear(in_features=2560, out_features=1024, bias=False)\n",
       "            (o_proj): Linear(in_features=2048, out_features=2560, bias=False)\n",
       "            (q_norm): Gemma3RMSNorm((256,), eps=1e-06)\n",
       "            (k_norm): Gemma3RMSNorm((256,), eps=1e-06)\n",
       "          )\n",
       "          (mlp): Gemma3MLP(\n",
       "            (gate_proj): Linear(in_features=2560, out_features=10240, bias=False)\n",
       "            (up_proj): Linear(in_features=2560, out_features=10240, bias=False)\n",
       "            (down_proj): Linear(in_features=10240, out_features=2560, bias=False)\n",
       "            (act_fn): PytorchGELUTanh()\n",
       "          )\n",
       "          (input_layernorm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
       "          (post_attention_layernorm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
       "          (pre_feedforward_layernorm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
       "          (post_feedforward_layernorm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
       "        )\n",
       "      )\n",
       "      (norm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
       "      (rotary_emb): Gemma3RotaryEmbedding()\n",
       "      (rotary_emb_local): Gemma3RotaryEmbedding()\n",
       "    )\n",
       "    (lm_head): Linear(in_features=2560, out_features=262208, bias=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoProcessor, Gemma3ForConditionalGeneration\n",
    "model = Gemma3ForConditionalGeneration.from_pretrained(model_path, device_map=\"auto\",\n",
    "                                             torch_dtype=torch.bfloat16, \n",
    "                                             trust_remote_code=True)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f3b0d02c-4256-4616-8902-9b6b794bb242",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.enable_input_require_grads() # 开启梯度检查点时，要执行该方法"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b6504777-7b98-4155-804e-c1a71d9ee127",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.bfloat16"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.dtype"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "309bd37b-6af6-42db-a480-3f99e2301eff",
   "metadata": {},
   "source": [
    "# LoRA "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a28da2cd-34ec-4aa3-aef0-138f13aa933b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import LoraConfig, TaskType, get_peft_model\n",
    "\n",
    "config = LoraConfig(\n",
    "    task_type=TaskType.CAUSAL_LM, \n",
    "    target_modules=[\"q_proj\", \"k_proj\",\"v_proj\", \"o_proj\"], # 可以自行添加更多微调的target_modules\n",
    "    inference_mode=False, # 训练模式\n",
    "    r=8,                  # Lora 秩\n",
    "    lora_alpha=32,        # Lora alaph，具体作用参见 Lora 原理\n",
    "    lora_dropout=0.1      # Dropout 比例\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1bc199e0-ae3e-4a08-81c2-0b963cea29cb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PeftModelForCausalLM(\n",
       "  (base_model): LoraModel(\n",
       "    (model): Gemma3ForConditionalGeneration(\n",
       "      (vision_tower): SiglipVisionModel(\n",
       "        (vision_model): SiglipVisionTransformer(\n",
       "          (embeddings): SiglipVisionEmbeddings(\n",
       "            (patch_embedding): Conv2d(3, 1152, kernel_size=(14, 14), stride=(14, 14), padding=valid)\n",
       "            (position_embedding): Embedding(4096, 1152)\n",
       "          )\n",
       "          (encoder): SiglipEncoder(\n",
       "            (layers): ModuleList(\n",
       "              (0-26): 27 x SiglipEncoderLayer(\n",
       "                (self_attn): SiglipSdpaAttention(\n",
       "                  (k_proj): lora.Linear(\n",
       "                    (base_layer): Linear(in_features=1152, out_features=1152, bias=True)\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=1152, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=1152, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                    (lora_magnitude_vector): ModuleDict()\n",
       "                  )\n",
       "                  (v_proj): lora.Linear(\n",
       "                    (base_layer): Linear(in_features=1152, out_features=1152, bias=True)\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=1152, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=1152, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                    (lora_magnitude_vector): ModuleDict()\n",
       "                  )\n",
       "                  (q_proj): lora.Linear(\n",
       "                    (base_layer): Linear(in_features=1152, out_features=1152, bias=True)\n",
       "                    (lora_dropout): ModuleDict(\n",
       "                      (default): Dropout(p=0.1, inplace=False)\n",
       "                    )\n",
       "                    (lora_A): ModuleDict(\n",
       "                      (default): Linear(in_features=1152, out_features=8, bias=False)\n",
       "                    )\n",
       "                    (lora_B): ModuleDict(\n",
       "                      (default): Linear(in_features=8, out_features=1152, bias=False)\n",
       "                    )\n",
       "                    (lora_embedding_A): ParameterDict()\n",
       "                    (lora_embedding_B): ParameterDict()\n",
       "                    (lora_magnitude_vector): ModuleDict()\n",
       "                  )\n",
       "                  (out_proj): Linear(in_features=1152, out_features=1152, bias=True)\n",
       "                )\n",
       "                (layer_norm1): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)\n",
       "                (mlp): SiglipMLP(\n",
       "                  (activation_fn): PytorchGELUTanh()\n",
       "                  (fc1): Linear(in_features=1152, out_features=4304, bias=True)\n",
       "                  (fc2): Linear(in_features=4304, out_features=1152, bias=True)\n",
       "                )\n",
       "                (layer_norm2): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "          (post_layernorm): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)\n",
       "        )\n",
       "      )\n",
       "      (multi_modal_projector): Gemma3MultiModalProjector(\n",
       "        (mm_soft_emb_norm): Gemma3RMSNorm((1152,), eps=1e-06)\n",
       "        (avg_pool): AvgPool2d(kernel_size=4, stride=4, padding=0)\n",
       "      )\n",
       "      (language_model): Gemma3ForCausalLM(\n",
       "        (model): Gemma3TextModel(\n",
       "          (embed_tokens): Gemma3TextScaledWordEmbedding(262208, 2560, padding_idx=0)\n",
       "          (layers): ModuleList(\n",
       "            (0-33): 34 x Gemma3DecoderLayer(\n",
       "              (self_attn): Gemma3Attention(\n",
       "                (q_proj): lora.Linear(\n",
       "                  (base_layer): Linear(in_features=2560, out_features=2048, bias=False)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.1, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=2560, out_features=8, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=8, out_features=2048, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                  (lora_magnitude_vector): ModuleDict()\n",
       "                )\n",
       "                (k_proj): lora.Linear(\n",
       "                  (base_layer): Linear(in_features=2560, out_features=1024, bias=False)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.1, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=2560, out_features=8, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=8, out_features=1024, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                  (lora_magnitude_vector): ModuleDict()\n",
       "                )\n",
       "                (v_proj): lora.Linear(\n",
       "                  (base_layer): Linear(in_features=2560, out_features=1024, bias=False)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.1, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=2560, out_features=8, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=8, out_features=1024, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                  (lora_magnitude_vector): ModuleDict()\n",
       "                )\n",
       "                (o_proj): lora.Linear(\n",
       "                  (base_layer): Linear(in_features=2048, out_features=2560, bias=False)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.1, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=2048, out_features=8, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=8, out_features=2560, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                  (lora_magnitude_vector): ModuleDict()\n",
       "                )\n",
       "                (q_norm): Gemma3RMSNorm((256,), eps=1e-06)\n",
       "                (k_norm): Gemma3RMSNorm((256,), eps=1e-06)\n",
       "              )\n",
       "              (mlp): Gemma3MLP(\n",
       "                (gate_proj): Linear(in_features=2560, out_features=10240, bias=False)\n",
       "                (up_proj): Linear(in_features=2560, out_features=10240, bias=False)\n",
       "                (down_proj): Linear(in_features=10240, out_features=2560, bias=False)\n",
       "                (act_fn): PytorchGELUTanh()\n",
       "              )\n",
       "              (input_layernorm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
       "              (post_attention_layernorm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
       "              (pre_feedforward_layernorm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
       "              (post_feedforward_layernorm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
       "            )\n",
       "          )\n",
       "          (norm): Gemma3RMSNorm((2560,), eps=1e-06)\n",
       "          (rotary_emb): Gemma3RotaryEmbedding()\n",
       "          (rotary_emb_local): Gemma3RotaryEmbedding()\n",
       "        )\n",
       "        (lm_head): Linear(in_features=2560, out_features=262208, bias=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = get_peft_model(model, config)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "4cc496b1-7e61-4a16-a3ea-cd210db48f45",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 5,949,440 || all params: 4,306,028,912 || trainable%: 0.1382\n"
     ]
    }
   ],
   "source": [
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "279416f1-7c17-4962-8cd6-e43302883778",
   "metadata": {},
   "source": [
    "# 配置训练参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d551d4cc-b8c9-46bb-9014-1ab1060361c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # 屏蔽令人讨厌的保存模型权重时UserWarning（Could not find a config file in model_path will assume that the vocabulary was not modified.）\n",
    "# 建议训练时打开\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\") \n",
    "\n",
    "output_dir=\"/root/autodl-tmp/LLM-Research/gemma-3-4b-it_lora_output\"\n",
    "\n",
    "args = TrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    per_device_train_batch_size=1,\n",
    "    gradient_accumulation_steps=4,\n",
    "    logging_steps=10,\n",
    "    num_train_epochs=3,\n",
    "    save_steps=100, \n",
    "    learning_rate=1e-4,\n",
    "    save_on_each_node=True,\n",
    "    gradient_checkpointing=True  # 开启梯度检查点，可以节省显存，加快训练速度，但会消耗更多内存， 对应model.enable_input_require_grads() \n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d19bbec3-3f64-4d80-ba59-6bc8c0694c06",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No label_names provided for model class `PeftModelForCausalLM`. Since `PeftModel` hides base models input arguments, if label_names is not given, label_names can't be set automatically within `Trainer`. Note that empty label_names list will be used instead.\n"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=tokenized_id,\n",
    "    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "81130465-f154-48c8-a4bb-4d750d29e89a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "It is strongly recommended to train Gemma3 models with the `eager` attention implementation instead of `sdpa`. Use `eager` with `AutoModelForCausalLM.from_pretrained('<path-to-checkpoint>', attn_implementation='eager')`.\n",
      "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='2796' max='2796' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [2796/2796 52:40, Epoch 2/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>27.133600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>20</td>\n",
       "      <td>21.347000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>30</td>\n",
       "      <td>18.589100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>40</td>\n",
       "      <td>18.701000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>50</td>\n",
       "      <td>16.447900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>60</td>\n",
       "      <td>15.465600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>70</td>\n",
       "      <td>15.476800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>80</td>\n",
       "      <td>14.833600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>90</td>\n",
       "      <td>14.242100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>13.849500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>110</td>\n",
       "      <td>15.143100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>120</td>\n",
       "      <td>13.934600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>130</td>\n",
       "      <td>13.579500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>140</td>\n",
       "      <td>12.739000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>150</td>\n",
       "      <td>14.071300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>160</td>\n",
       "      <td>13.429800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>170</td>\n",
       "      <td>13.409500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>180</td>\n",
       "      <td>14.614900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>190</td>\n",
       "      <td>13.298300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>13.242100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>210</td>\n",
       "      <td>13.760000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>220</td>\n",
       "      <td>13.903900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>230</td>\n",
       "      <td>14.041100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>240</td>\n",
       "      <td>13.384900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>250</td>\n",
       "      <td>14.652700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>260</td>\n",
       "      <td>12.969400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>270</td>\n",
       "      <td>12.975600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>280</td>\n",
       "      <td>13.404200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>290</td>\n",
       "      <td>13.863100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>300</td>\n",
       "      <td>14.716900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>310</td>\n",
       "      <td>14.040800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>320</td>\n",
       "      <td>12.587900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>330</td>\n",
       "      <td>12.495100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>340</td>\n",
       "      <td>14.416900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>350</td>\n",
       "      <td>13.454700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>360</td>\n",
       "      <td>14.047200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>370</td>\n",
       "      <td>13.050200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>380</td>\n",
       "      <td>12.059400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>390</td>\n",
       "      <td>12.534500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>400</td>\n",
       "      <td>14.332300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>410</td>\n",
       "      <td>13.280900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>420</td>\n",
       "      <td>14.451200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>430</td>\n",
       "      <td>13.236300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>440</td>\n",
       "      <td>13.571700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>450</td>\n",
       "      <td>13.524600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>460</td>\n",
       "      <td>13.190800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>470</td>\n",
       "      <td>13.185000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>480</td>\n",
       "      <td>12.262100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>490</td>\n",
       "      <td>12.489400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>12.732000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>510</td>\n",
       "      <td>13.114500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>520</td>\n",
       "      <td>14.342300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>530</td>\n",
       "      <td>14.097600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>540</td>\n",
       "      <td>12.669000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>550</td>\n",
       "      <td>12.153500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>560</td>\n",
       "      <td>13.363600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>570</td>\n",
       "      <td>13.394200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>580</td>\n",
       "      <td>12.186100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>590</td>\n",
       "      <td>13.366800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>600</td>\n",
       "      <td>13.428600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>610</td>\n",
       "      <td>13.010300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>620</td>\n",
       "      <td>14.218500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>630</td>\n",
       "      <td>12.775500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>640</td>\n",
       "      <td>13.784900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>650</td>\n",
       "      <td>12.800500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>660</td>\n",
       "      <td>13.242400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>670</td>\n",
       "      <td>12.458800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>680</td>\n",
       "      <td>13.079300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>690</td>\n",
       "      <td>12.395200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>700</td>\n",
       "      <td>11.744200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>710</td>\n",
       "      <td>12.563100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>720</td>\n",
       "      <td>13.779400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>730</td>\n",
       "      <td>12.900400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>740</td>\n",
       "      <td>12.184500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>750</td>\n",
       "      <td>11.759400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>760</td>\n",
       "      <td>11.909800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>770</td>\n",
       "      <td>13.100600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>780</td>\n",
       "      <td>13.054700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>790</td>\n",
       "      <td>11.948200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>800</td>\n",
       "      <td>13.018200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>810</td>\n",
       "      <td>14.420800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>820</td>\n",
       "      <td>12.306400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>830</td>\n",
       "      <td>12.885100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>840</td>\n",
       "      <td>11.820400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>850</td>\n",
       "      <td>11.164000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>860</td>\n",
       "      <td>13.010700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>870</td>\n",
       "      <td>13.241200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>880</td>\n",
       "      <td>13.301700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>890</td>\n",
       "      <td>12.702000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>900</td>\n",
       "      <td>13.442800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>910</td>\n",
       "      <td>12.595400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>920</td>\n",
       "      <td>13.190700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>930</td>\n",
       "      <td>13.797800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>940</td>\n",
       "      <td>11.768800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>950</td>\n",
       "      <td>12.767500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>960</td>\n",
       "      <td>12.375300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>970</td>\n",
       "      <td>11.610600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>980</td>\n",
       "      <td>11.844600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>990</td>\n",
       "      <td>11.325000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>12.058900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1010</td>\n",
       "      <td>10.899700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1020</td>\n",
       "      <td>11.994500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1030</td>\n",
       "      <td>11.704900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1040</td>\n",
       "      <td>11.676200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1050</td>\n",
       "      <td>10.593400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1060</td>\n",
       "      <td>11.973100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1070</td>\n",
       "      <td>12.000600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1080</td>\n",
       "      <td>11.815400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1090</td>\n",
       "      <td>11.436800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1100</td>\n",
       "      <td>12.153700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1110</td>\n",
       "      <td>12.267900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1120</td>\n",
       "      <td>11.660400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1130</td>\n",
       "      <td>11.483900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1140</td>\n",
       "      <td>13.132900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1150</td>\n",
       "      <td>13.530200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1160</td>\n",
       "      <td>10.713300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1170</td>\n",
       "      <td>10.883700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1180</td>\n",
       "      <td>11.058600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1190</td>\n",
       "      <td>12.374100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1200</td>\n",
       "      <td>12.718100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1210</td>\n",
       "      <td>12.706600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1220</td>\n",
       "      <td>12.324000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1230</td>\n",
       "      <td>10.218300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1240</td>\n",
       "      <td>11.259100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1250</td>\n",
       "      <td>12.191900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1260</td>\n",
       "      <td>10.984400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1270</td>\n",
       "      <td>12.198800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1280</td>\n",
       "      <td>11.847800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1290</td>\n",
       "      <td>12.257800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1300</td>\n",
       "      <td>11.900000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1310</td>\n",
       "      <td>11.586500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1320</td>\n",
       "      <td>12.289600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1330</td>\n",
       "      <td>11.941100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1340</td>\n",
       "      <td>10.703200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1350</td>\n",
       "      <td>12.183900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1360</td>\n",
       "      <td>11.044900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1370</td>\n",
       "      <td>11.332500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1380</td>\n",
       "      <td>13.122500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1390</td>\n",
       "      <td>12.043300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1400</td>\n",
       "      <td>11.035300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1410</td>\n",
       "      <td>11.430100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1420</td>\n",
       "      <td>11.944500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1430</td>\n",
       "      <td>11.218100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1440</td>\n",
       "      <td>12.226100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1450</td>\n",
       "      <td>11.405300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1460</td>\n",
       "      <td>12.749600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1470</td>\n",
       "      <td>11.846700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1480</td>\n",
       "      <td>11.604900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1490</td>\n",
       "      <td>12.057800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1500</td>\n",
       "      <td>11.820400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1510</td>\n",
       "      <td>12.523700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1520</td>\n",
       "      <td>11.595200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1530</td>\n",
       "      <td>11.262300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1540</td>\n",
       "      <td>12.352900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1550</td>\n",
       "      <td>11.983900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1560</td>\n",
       "      <td>12.054000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1570</td>\n",
       "      <td>10.725700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1580</td>\n",
       "      <td>10.493100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1590</td>\n",
       "      <td>11.607800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1600</td>\n",
       "      <td>11.746800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1610</td>\n",
       "      <td>11.440200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1620</td>\n",
       "      <td>12.634300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1630</td>\n",
       "      <td>12.863500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1640</td>\n",
       "      <td>12.078900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1650</td>\n",
       "      <td>12.261400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1660</td>\n",
       "      <td>12.390000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1670</td>\n",
       "      <td>11.103300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1680</td>\n",
       "      <td>11.621900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1690</td>\n",
       "      <td>12.873800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1700</td>\n",
       "      <td>11.474200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1710</td>\n",
       "      <td>10.678300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1720</td>\n",
       "      <td>11.178000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1730</td>\n",
       "      <td>11.302600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1740</td>\n",
       "      <td>10.563000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1750</td>\n",
       "      <td>11.708500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1760</td>\n",
       "      <td>12.005800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1770</td>\n",
       "      <td>10.724100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1780</td>\n",
       "      <td>10.598100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1790</td>\n",
       "      <td>11.424600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1800</td>\n",
       "      <td>11.857900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1810</td>\n",
       "      <td>12.254100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1820</td>\n",
       "      <td>10.849400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1830</td>\n",
       "      <td>11.654700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1840</td>\n",
       "      <td>11.842200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1850</td>\n",
       "      <td>11.841300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1860</td>\n",
       "      <td>11.610000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1870</td>\n",
       "      <td>9.994800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1880</td>\n",
       "      <td>11.234700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1890</td>\n",
       "      <td>9.631100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1900</td>\n",
       "      <td>9.725600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1910</td>\n",
       "      <td>11.134200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1920</td>\n",
       "      <td>9.617600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1930</td>\n",
       "      <td>9.195400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1940</td>\n",
       "      <td>11.525700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1950</td>\n",
       "      <td>10.249100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1960</td>\n",
       "      <td>11.692700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1970</td>\n",
       "      <td>10.598200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1980</td>\n",
       "      <td>9.821100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1990</td>\n",
       "      <td>10.650600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2000</td>\n",
       "      <td>10.869800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2010</td>\n",
       "      <td>11.417500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2020</td>\n",
       "      <td>9.355600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2030</td>\n",
       "      <td>10.887300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2040</td>\n",
       "      <td>10.432100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2050</td>\n",
       "      <td>11.677800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2060</td>\n",
       "      <td>10.629700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2070</td>\n",
       "      <td>10.937500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2080</td>\n",
       "      <td>9.539100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2090</td>\n",
       "      <td>10.302700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2100</td>\n",
       "      <td>10.111600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2110</td>\n",
       "      <td>9.845500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2120</td>\n",
       "      <td>10.648300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2130</td>\n",
       "      <td>10.955500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2140</td>\n",
       "      <td>10.503900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2150</td>\n",
       "      <td>11.879500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2160</td>\n",
       "      <td>11.138200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2170</td>\n",
       "      <td>11.775000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2180</td>\n",
       "      <td>10.980700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2190</td>\n",
       "      <td>10.989400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2200</td>\n",
       "      <td>10.217800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2210</td>\n",
       "      <td>11.486300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2220</td>\n",
       "      <td>11.496200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2230</td>\n",
       "      <td>10.566100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2240</td>\n",
       "      <td>10.561700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2250</td>\n",
       "      <td>10.774600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2260</td>\n",
       "      <td>11.621200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2270</td>\n",
       "      <td>9.948400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2280</td>\n",
       "      <td>10.848400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2290</td>\n",
       "      <td>10.021900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2300</td>\n",
       "      <td>11.530900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2310</td>\n",
       "      <td>9.752300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2320</td>\n",
       "      <td>9.879400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2330</td>\n",
       "      <td>10.337600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2340</td>\n",
       "      <td>8.862600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2350</td>\n",
       "      <td>11.082500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2360</td>\n",
       "      <td>11.376600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2370</td>\n",
       "      <td>10.513500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2380</td>\n",
       "      <td>10.111700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2390</td>\n",
       "      <td>11.432600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2400</td>\n",
       "      <td>10.896400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2410</td>\n",
       "      <td>10.551400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2420</td>\n",
       "      <td>12.013800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2430</td>\n",
       "      <td>10.691200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2440</td>\n",
       "      <td>11.052300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2450</td>\n",
       "      <td>9.853400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2460</td>\n",
       "      <td>11.384000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2470</td>\n",
       "      <td>10.882000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2480</td>\n",
       "      <td>10.998500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2490</td>\n",
       "      <td>9.820900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2500</td>\n",
       "      <td>11.692100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2510</td>\n",
       "      <td>11.007100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2520</td>\n",
       "      <td>10.614900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2530</td>\n",
       "      <td>10.080400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2540</td>\n",
       "      <td>10.767100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2550</td>\n",
       "      <td>10.712200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2560</td>\n",
       "      <td>10.911600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2570</td>\n",
       "      <td>10.830000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2580</td>\n",
       "      <td>10.686100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2590</td>\n",
       "      <td>10.838800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2600</td>\n",
       "      <td>10.170300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2610</td>\n",
       "      <td>11.887700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2620</td>\n",
       "      <td>11.243100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2630</td>\n",
       "      <td>9.874100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2640</td>\n",
       "      <td>11.052900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2650</td>\n",
       "      <td>11.489300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2660</td>\n",
       "      <td>12.466200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2670</td>\n",
       "      <td>11.467800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2680</td>\n",
       "      <td>9.634000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2690</td>\n",
       "      <td>9.034700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2700</td>\n",
       "      <td>10.596800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2710</td>\n",
       "      <td>9.692200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2720</td>\n",
       "      <td>10.336700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2730</td>\n",
       "      <td>10.506600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2740</td>\n",
       "      <td>10.627500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2750</td>\n",
       "      <td>10.666700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2760</td>\n",
       "      <td>9.632100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2770</td>\n",
       "      <td>10.286400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2780</td>\n",
       "      <td>9.931700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2790</td>\n",
       "      <td>11.270800</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=2796, training_loss=12.025420244841786, metrics={'train_runtime': 3162.7596, 'train_samples_per_second': 3.537, 'train_steps_per_second': 0.884, 'total_flos': 2.321808952845744e+16, 'train_loss': 12.025420244841786, 'epoch': 2.997586484312148})"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d725d46-6dc3-428b-870b-5275a0838457",
   "metadata": {},
   "source": [
    "# 合并加载模型\n",
    "\n",
    "> 这里推荐大家在**训练结束重启一下notebook, 释放微调占用的GPU显存**, 否则容易出现如下警告\"Some parameters are on the meta device because they were offloaded to the cpu.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0b484a15-a300-4be8-ab08-8722c1dc3b50",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f177d8f03bf4de8bfd9fba99946e991",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "prompt:  你是谁？\n",
      "system_prompt:  现在你要扮演皇帝身边的女人--甄嬛\n",
      "output:  我是甄嬛，家父是大理寺少卿甄远道。\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "import torch\n",
    "from peft import PeftModel\n",
    "\n",
    "model_path = '/root/autodl-tmp/LLM-Research/gemma-3-4b-it'\n",
    "lora_path = '/root/autodl-tmp/LLM-Research/gemma-3-4b-it_lora_output/checkpoint-2790' # 这里改成 LoRA 输出对应 checkpoint 地址和最终的 epoch 数值 2796\n",
    "\n",
    "# 加载tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)\n",
    "\n",
    "# 加载模型\n",
    "model = AutoModelForCausalLM.from_pretrained(model_path, \n",
    "                                             device_map=\"auto\",\n",
    "                                             torch_dtype=torch.bfloat16, \n",
    "                                             trust_remote_code=True).eval()\n",
    "\n",
    "# 加载lora权重\n",
    "model = PeftModel.from_pretrained(model, model_id=lora_path)\n",
    "\n",
    "prompt = \"你是谁？\"\n",
    "system_prompt = \"现在你要扮演皇帝身边的女人--甄嬛\"\n",
    "print(\"prompt: \", prompt)\n",
    "print(\"system_prompt: \", system_prompt)\n",
    "\n",
    "inputs = tokenizer.apply_chat_template([{\"role\": \"system\", \"content\": system_prompt},\n",
    "                                        {\"role\": \"user\", \"content\": prompt}],\n",
    "                                       add_generation_prompt=True,\n",
    "                                       tokenize=True,\n",
    "                                       return_tensors=\"pt\",\n",
    "                                       return_dict=True\n",
    "                                       ).to(model.device)  # 将 inputs 移动到模型所在的设备，确保设备一致性\n",
    "\n",
    "\n",
    "gen_kwargs = {\"max_length\": 2500, \"do_sample\": True, \"top_k\": 1}\n",
    "with torch.no_grad():\n",
    "    outputs = model.generate(**inputs, **gen_kwargs)\n",
    "    outputs = outputs[:, inputs['input_ids'].shape[1]:]\n",
    "    print(\"output: \", tokenizer.decode(outputs[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "441973d4-6dd0-40b4-bb79-cfd53752a742",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
